package rerac.feature;

import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;

import rerac.protos.Corpus.Document;
import rerac.protos.Corpus.Document.Annotation;
import rerac.protos.Corpus.Document.Token;
import rerac.protos.Corpus.Document.AnnotationType;

import com.google.common.collect.HashMultimap;
import com.google.common.collect.Multimap;

/**
 * This holds a dependency tree, typically constructed from a parsed document.
 * At the moment used to get a string feature for pairs of tokens from their
 * shortest path in the tree.
 * 
 * @author Benjamin Roth
 *
 */
public class DependencyTree {
  
  /**
   * An edge pointing to another node in the tree.
   * 
   * @author Benjamin Roth
   *
   */
  class Edge {
    // The node index pointed to.
    private int targetIndex;
    // The edge label.
    private String label;
    
    /**
     * 
     * @param l a String, the Edge label.
     * @param target the int index of the node pointed to.
     */
    Edge(String l, int target) {
      label = l;
      targetIndex = target;
    }
    
    /**
     * 
     * @return an int, the index og the node pointed to.
     */
    int getTarget() {
      return targetIndex;
    }
    /**
     * 
     * @return a String, the label of the edge.
     */
    String getLabel() {
      return label;
    }
  }

  // This provides a mapping from any node index that has outgoing edges to the 
  // edges going out from that node.
  Multimap<Integer, Edge> index2edges = HashMultimap.create();
  // The number of tokens = potential nodes.
  int numTokens;
  
  /**
   * This builds up a parse tree from a parsed document.
   * 
   * @param doc a Document in protobuf format.
   */
  DependencyTree(Document doc) {
    for (int i = 0; i < doc.getTokenCount(); ++i) {
      Token t = doc.getToken(i);
      for (Annotation a : t.getAnnotationList()) {
        if (a.getType() == AnnotationType.SYNTAX_HEAD && 
            a.hasTargetToken()) {
          String label = a.hasText() ? a.getText() : "";
          Edge e = new Edge(label, a.getTargetToken());
          index2edges.put(i, e);
        }
      }
    }
    numTokens = doc.getTokenCount();
  }
  
  /**
   * 
   * This returns a String representation of the shortest path between two 
   * nodes. The string representation is made up from the labels of the edges
   * and indicates the directionality of the individual edges.
   * 
   * For example if the shortes path from token_2 to token_7 goes through 
   * token_5 and the following dependencies exist 
   * 
   * child edge_label parent
   * -----------------------
   * token_2 SUBJ token_5
   * token_7 OBJ token_5
   * 
   * Then the String pathString(2,7) returns
   * 
   * SUBJ-><-OBJ
   * 
   * Note, that pathString(7,2) would return
   * 
   * OBJ-><-SUBJ
   * 
   * @param indexA index of a token in the sentence
   * @param indexB index of another token in the sentence
   * @return The string representation of the shortest path connecting indexA
   * with indexB, the empty String if the is no such path.
   */
  public String pathString(int indexA, int indexB) {
    List<Integer> path = shortestPath(indexA, indexB);
    StringBuffer sb = new StringBuffer();
    
    int lastIndex = indexA;
    for (int i = 1; i < path.size(); ++i) {
      int currIndex = path.get(i);
      if (i > 1) {
        sb.append(" ");
      }
      boolean upEdge = false;
      for (Edge e : index2edges.get(lastIndex)) {
        if (e.getTarget() == currIndex) {
          sb.append(e.getLabel());
          sb.append("->");
          upEdge=true;
          break;
        }
      }
      if (!upEdge) {
        for (Edge e : index2edges.get(currIndex)) {
          if (e.getTarget() == lastIndex) {
            sb.append("<-");
            sb.append(e.getLabel());
            break;
          }
        }        
      }
      lastIndex = currIndex;
    }
    
    return sb.toString();
  }
  
  /**
   * This finds the shortest path from A to B. This is done by breadth first 
   * search on the transition matrix of the undirected graph of the dependency
   * tree (i.e. egdes are followed in both directions).
   * 
   * @param indexA index of a token in the sentence.
   * @param indexB index of another token in the sentence.
   * @return The indices of the nodes the shortest path traverses.
   */
  private List<Integer> shortestPath(int indexA, int indexB) {
    Multimap<Integer, Integer> transitions = 
      HashMultimap.create(numTokens,index2edges.values().size() * 2);
    boolean[] blocked = new boolean[numTokens];
    int[] comingFrom = new int[numTokens];
    
    for (int i1 : index2edges.keySet()) {
      for (Edge e : index2edges.get(i1)) {
        int i2 = e.getTarget();
        transitions.put(i1, i2);
        transitions.put(i2, i1);
      }
    }
    
    LinkedList<Integer> toVisit = new LinkedList<Integer>();
    toVisit.add(indexA);
    
    boolean foundPath = false;
    while (toVisit.size() > 0 && !foundPath) {
      int currNode = toVisit.poll();
      for (int nextNode : transitions.get(currNode)) {
        if (-1 == nextNode) {
          continue;
        }
        if (!blocked[nextNode]) {
          comingFrom[nextNode] = currNode;
          blocked[nextNode] = true;
          toVisit.add(nextNode);
          if (nextNode == indexB) {
            foundPath = true;
            break;
          }
        }
      }
    }
    
    List<Integer> path = new LinkedList<Integer>();
    
    if (foundPath) {
      for (int nextNode = indexB; nextNode != indexA; 
          nextNode = comingFrom[nextNode]) {
        path.add(0, nextNode);
      }
      path.add(0, indexA);
    }
    
    return new ArrayList<Integer>(path);
  }
}
